#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/timing.h>
#include <src/arena.h>
#include <src/coeff_value.h>
#include <src/accumulator.h>
#include <src/accumulator_energy.h>
#include <src/accumulator_properties.h>
#include <src/accumulator_gradient.h>
#include <src/accumulator_stoch_recon.h>
#include <src/markov_chain.h>
#include <src/stochastic_optimization.h>

//-------------------------------------------------------------------------------
// pcps::stochastic_iterations-- Iteratively performs stochastic sampling of a
//                               wavefunction.  This function can be used to
//                               optimize the wavefunction, or simply to compute
//                               its properties.
//
//                               The behavior is governed by which accumulators
//                               are provided, how many iterations are requested,
//                               and the sample and burn in lengths.
//
//                               print_freq and dump_freq dictate which
//                               iterations should print information and dump
//                               the wavefunction to disk.
//-------------------------------------------------------------------------------

template <class S, class PART, class REF>
void pcps::stochastic_iterations(const pcps::Input & userinp,
                                 const pcps::Hamiltonian<S> & ham,
                                 pcps::Wavefunction<S> & wfn,
                                 std::vector<boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > > & accumulators,
                                 const int max_iter,
                                 const int print_freq,
                                 const int dump_freq,
                                 const long int sample_length,   // number of samples to take on EACH thread
                                 const long int burn_in_length)  // number burn in steps for each thread's markov chain
{

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // get the number of threads
  const int nthread = userinp.nthreads();

  // get the number of configurations on a site
  const int configs_per_site = userinp.configs_per_site();

  // get the number of sites
  const int nsites = userinp.nsites();

  // check that we do not have more than one of any accumulator
  for (int i = 0; i < accumulators.size(); i++)
    for (int j = i+1; j < accumulators.size(); j++)
      if ( typeid(*accumulators.at(i)) == typeid(*accumulators.at(j)) )
        throw pcps::Exception( (boost::format("Error in pcps::stochastic_iterations.  Cannot have more than one accumulator of a given type.")).str() );

  // ensure that each accumulator has the correct sample length
  for (int i = 0; i < accumulators.size(); i++)
    if ( accumulators.at(i)->sample_length() != sample_length )
      throw pcps::Exception( (boost::format("Error in pcps::stochastic_iterations.  Accumulator %i has the wrong sample length.") % i).str() );

  // Create a markov chain object, which spawns a markov chain on each thread and then waits for instructions.
  // For each accumulator, each thread will now own a thread accumulator object associated with the corresponding process accumulator.
  pcps::MarkovChain<S, PART, REF> markov(userinp, ham, wfn, accumulators);

  // burn in each thread's markov chain
  markov.burn_in(burn_in_length);

  // iteratevely perform the accumulators' actions
  for (int iter = 0; iter < max_iter; iter++) {

    // reset the thread accumulators
    markov.reset_accumulators();

    // reset the process accumulators
    for (int i = 0; i < accumulators.size(); i++)
      accumulators.at(i)->reset();

    // reset each thread's coefficient value function object to reflect the new wavefunction
    markov.reset_cvf();

    // have each thread sample the markov chain
    markov.sample(sample_length);

    // use the thread accumulators to compute the wavefunction update (some accumulators, like the energy, do nothing here)
    markov.compute_update();

    // use the process accumulators to update the wavefunction (some accumulators, like the energy, do nothing here)
    for (int i = 0; i < accumulators.size(); i++)
      accumulators.at(i)->update_wfn(userinp, wfn);

    // scale the wavefunction's variables at each dump iteration
    if ((iter+1) % dump_freq == 0)
      wfn.scale_variables();

    // save the wavefunction to disk at each dump iteration
    if (myrank == 0 && (iter+1) % dump_freq == 0)
      wfn.dump_binary(userinp);

    // print information at each print iteration
    if ((iter+1) % print_freq == 0) {
      if (myrank == 0)
        std::cout << boost::format("stochastic iteration %22i        elapsed time = %17.6f sec")
                     % (iter+1)
                     % pcps::get_overall_timer().elapsed_seconds()
                  << std::endl;
      for (int i = 0; i < accumulators.size(); i++)
        accumulators.at(i)->print_iter_info(userinp, wfn, iter, accumulators);
      if (myrank == 0)
        std::cout << std::endl;
      if (myrank == 0 && userinp.verbose_print())
        wfn.print(userinp);
    }

  }

}

//-------------------------------------------------------------------------------
// pcps::stochastic_opt_detail -- Optimizes the wavefunction using one of a
//                                number of algorithms that rely on stochastic
//                                sampling.
//
//                                This function chooses which optimization to
//                                perform, creates the necessary accumulators,
//                                and sends the necessary instructions to the
//                                stochastic_iterations function.
//-------------------------------------------------------------------------------

template <class S, class PART, class REF>
void pcps::stochastic_opt_detail(const pcps::Input & userinp, const pcps::Hamiltonian<S> & ham, pcps::Wavefunction<S> & wfn) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // divide the sample length among the different threads
  const long int sample_length = 1 + userinp.sample_length() / ( nproc * userinp.nthreads() );

  // define a short name for a shared pointer to a process accumulator
  typedef boost::shared_ptr<pcps::ProcessAccum<S, PART, REF> > ProcAccPtr;

  // initialize a vector for the accumulators
  std::vector<ProcAccPtr> accumulators;

  // the energy accumulator comes first for all optimization methods
  { ProcAccPtr temp( new pcps::EnergyProcessAccum<S, PART, REF>(userinp, ham, wfn, sample_length) ); accumulators.push_back(temp); }

  // if requested, include an accumulator to compute correlation functions
  if (userinp.properties())
    { ProcAccPtr temp( new pcps::PropertiesProcessAccum<S, PART, REF>(userinp, wfn, sample_length) ); accumulators.push_back(temp); }

  // variational steepest descent
  if (userinp.optimization() == pcps::keyword::Opt_Steepest_Descent)
    { ProcAccPtr temp( new pcps::GradientProcessAccum<S, PART, REF>(userinp, wfn, sample_length) ); accumulators.push_back(temp); }

  // stochastic reconfiguration
  else if ( userinp.optimization() == pcps::keyword::Opt_Stochastic_Imaginary_Time ||
            userinp.optimization() == pcps::keyword::Opt_Stochastic_Real_Time         )
    { ProcAccPtr temp( new pcps::StochReconProcessAccum<S, PART, REF>(userinp, wfn, sample_length) ); accumulators.push_back(temp); }

  // unknown optimization
  else
    throw pcps::Exception( (boost::format("\"%s\" is not an allowed stochastic optimization technique") % userinp.optimization()).str() );

  // run the iterative optimization
  pcps::stochastic_iterations<S, PART, REF>(userinp, ham, wfn, accumulators, userinp.max_macro_loop(),
                                                 userinp.print_freq(), userinp.dump_freq(), sample_length, userinp.burn_in_length());

  // scale the wavefunction's variables
  wfn.scale_variables();

  // save the wavefunction to disk
  if (myrank == 0)
    wfn.dump_binary(userinp);

}

//-------------------------------------------------------------------------------
// pcps::stochastic_optimization -- Optimizes the wavefunction using one of a
//                                  number of algorithms that rely on stochastic
//                                  sampling.
//
//                                  This function chooses which particle type and
//                                  reference type we are using and calls the
//                                  appropriate version of stochastic_opt_detail.
//-------------------------------------------------------------------------------

template <class S>
void pcps::stochastic_optimization(const pcps::Input & userinp, const pcps::Hamiltonian<S> & ham, pcps::Wavefunction<S> & wfn) {

  // spins with the uniform reference
  if (userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_Uniform)
    pcps::stochastic_opt_detail<S, pcps::Part_Spin, pcps::Ref_Uniform>(userinp, ham, wfn);

  // spins with the AGP reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_AGP )
    pcps::stochastic_opt_detail<S, pcps::Part_Spin, pcps::Ref_AGP>(userinp, ham, wfn);

  // spins with the pairing reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Spin && userinp.ref_type() == pcps::keyword::Ref_Pairing )
    pcps::stochastic_opt_detail<S, pcps::Part_Spin, pcps::Ref_Pairing>(userinp, ham, wfn);

  // fermions with the uniform reference
  else if (userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Uniform)
    pcps::stochastic_opt_detail<S, pcps::Part_Fermion, pcps::Ref_Uniform>(userinp, ham, wfn);

  // fermions with the restricted Hartree Fock or AGP reference
  else if (   userinp.particle_type() == pcps::keyword::Part_Fermion
           && ( userinp.ref_type() == pcps::keyword::Ref_RHF || userinp.ref_type() == pcps::keyword::Ref_AGP ) )
    pcps::stochastic_opt_detail<S, pcps::Part_Fermion, pcps::Ref_AGP>(userinp, ham, wfn);

  // fermions with the pairing reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Pairing )
    pcps::stochastic_opt_detail<S, pcps::Part_Fermion, pcps::Ref_Pairing>(userinp, ham, wfn);

  // fermions with the pfaffian reference
  else if ( userinp.particle_type() == pcps::keyword::Part_Fermion && userinp.ref_type() == pcps::keyword::Ref_Pfaffian )
    pcps::stochastic_opt_detail<S, pcps::Part_Fermion, pcps::Ref_Pfaffian>(userinp, ham, wfn);

  // unsupported combination of particle and reference type
  else
    throw pcps::Exception(
      (boost::format("Error in pcps::stochastic_optimization.  The combination of particle type \"%s\" and reference type \"%s\" is not supported.")
        % userinp.particle_type()
        % userinp.ref_type()
      ).str()
    );

}

// instantiate the template for real and complex numbers
template void pcps::stochastic_optimization(const pcps::Input &,
                                            const pcps::Hamiltonian<double> &,
                                            pcps::Wavefunction<double> &);
template void pcps::stochastic_optimization(const pcps::Input &,
                                            const pcps::Hamiltonian<std::complex<double> > &,
                                            pcps::Wavefunction<std::complex<double> > &);
